import { useCodeStyle } from "@renderer/hooks/use-code-style";
import { codeToHtmlCached } from "@renderer/lib/shiki-highlighter";
import { debounce } from "es-toolkit";
import { useCallback, useMemo, useRef, useState } from "react";
import type { ThemedToken } from "shiki/core";

export interface UseCodeHighlightOptions {
  rawLines: string[];
  language: string;
}

export function useCodeHighlight({
  rawLines,
  language,
}: UseCodeHighlightOptions) {
  const { isShikiThemeDark } = useCodeStyle();

  const themeMode = isShikiThemeDark ? "dark" : "light";
  const [tokenLines, setTokenLines] = useState<(ThemedToken[] | undefined)[]>(
    [],
  );
  const lastProcessedRef = useRef<number>(0);
  const processingRef = useRef<boolean>(false);

  const parseShikiHtml = useCallback(
    (html: string): ThemedToken[][] => {
      const lineTokens: ThemedToken[][] = [];

      for (let i = 0; i < rawLines.length; i++) {
        lineTokens[i] = [];
      }

      const tempDiv = document.createElement("div");
      tempDiv.innerHTML = html;

      const codeElement = tempDiv.querySelector("pre code");
      if (!codeElement) {
        rawLines.forEach((line, index) => {
          lineTokens[index] = [
            {
              content: line,
              offset: 0,
            },
          ];
        });
        return lineTokens;
      }

      const lineElements = codeElement.querySelectorAll(".line");

      if (lineElements.length > 0) {
        lineElements.forEach((lineElement, lineIndex) => {
          if (lineIndex >= rawLines.length) return;

          const tokens: ThemedToken[] = [];
          const spans = lineElement.querySelectorAll("span");
          let offset = 0;

          if (spans.length === 0) {
            const textContent = lineElement.textContent || "";
            tokens.push({
              content: textContent,
              offset: 0,
            });
          } else {
            spans.forEach((span) => {
              const content = span.textContent || "";
              const style = span.getAttribute("style") || "";

              const colorMatch = style.match(/color:\s*([^;]+)/);
              const bgColorMatch = style.match(/background-color:\s*([^;]+)/);

              tokens.push({
                content,
                offset,
                color: colorMatch?.[1]?.trim(),
                bgColor: bgColorMatch?.[1]?.trim(),
              });
              offset += content.length;
            });
          }

          lineTokens[lineIndex] = tokens;
        });

        for (let i = lineElements.length; i < rawLines.length; i++) {
          lineTokens[i] = [
            {
              content: rawLines[i] || "",
              offset: 0,
            },
          ];
        }
      } else {
        const fullText = codeElement.textContent || "";
        const shikiLines = fullText.split("\n");

        rawLines.forEach((_, index) => {
          const shikiLine = index < shikiLines.length ? shikiLines[index] : "";

          lineTokens[index] = [
            {
              content: shikiLine,
              offset: 0,
            },
          ];
        });
      }

      return lineTokens;
    },
    [rawLines],
  );

  const highlightLines = useCallback(
    async (endLine: number) => {
      if (processingRef.current) return;

      const linesToProcess = Math.min(endLine, rawLines.length);
      if (linesToProcess <= lastProcessedRef.current) return;

      processingRef.current = true;

      try {
        const content = rawLines.slice(0, linesToProcess).join("\n");

        const html = await codeToHtmlCached(content, language, themeMode);

        if (html) {
          const allLineTokens = parseShikiHtml(html);

          setTokenLines((prev) => {
            const newTokenLines = [...prev];
            allLineTokens.forEach((tokens, index) => {
              if (index < rawLines.length) {
                newTokenLines[index] = tokens;
              }
            });
            return newTokenLines;
          });

          lastProcessedRef.current = linesToProcess;
        }
      } catch (_error) {
      } finally {
        processingRef.current = false;
      }
    },
    [rawLines, language, themeMode, parseShikiHtml],
  );

  const debouncedHighlightLines = useMemo(
    () => debounce(highlightLines, 300),
    [highlightLines],
  );

  const resetTokens = useCallback(() => {
    setTokenLines([]);
    lastProcessedRef.current = 0;
  }, []);

  return {
    tokenLines,
    highlightLines: debouncedHighlightLines,
    resetTokens,
  };
}
